import { create } from 'zustand'
import { persist } from 'zustand/middleware'

export enum gptModule {
  gtp_4 = 1,
  gtp_3_5 = 2,
  xf_model = 3,
}

export interface module {
  gtp_4: gptModule.gtp_4
  gtp_3_5: gptModule.gtp_3_5
  xf_model: gptModule.xf_model
  [key: string]: any
}

type ModuleKey = keyof module

interface moduleItem {
  id: string | number
  name: ModuleKey
}

export interface ChatItem {
  id: string | number
  user_id: string
  sender_type: 'user' | 'chat'
  message: string
  temporary?: string | number
}

interface ReceiveMsg {
  sender_type: 'user' | 'chat'
  temporary?: string | number
  data?: {
    message: string
    [key: string]: any
  }
  [key: string]: any
}

interface Chat {
  module: ModuleKey
  moduleList: moduleItem[]
  list: ChatItem[]
  isSend: boolean
  setModule: (module: ModuleKey) => void
  setModuleList: (moduleList: moduleItem[]) => void
  setMsgList: (list: ChatItem[]) => void
  getMsgList: () => ChatItem[]
  clear: () => void
  receive: (msg: ReceiveMsg) => void
  sendMsg: (msg: string) => void
}

export const useStore = create<Chat>()(
  // 可选：如果不需要持久化可移除 `persist`
  persist(
    (set, get) => ({
      module: 'gtp_3_5',
      moduleList: [
        { id: gptModule.gtp_3_5, name: 'gtp_3_5' },
        { id: gptModule.gtp_4, name: 'gtp_4' },
        { id: gptModule.xf_model, name: 'xf_model' },
      ],
      list: [],
      isSend: false,

      setModule: (id: ModuleKey) => {
        const found = get().moduleList.find((item) => item.name === id)
        if (!found) return set({ module: 'gtp_3_5' }) // fallback
        set({ module: id })
      },

      setModuleList: (moduleList: moduleItem[]) => {
        set({ moduleList })
      },

      setMsgList: (list: ChatItem[]) => {
        set({ list })
      },

      getMsgList: () => {
        return get().list
      },

      sendMsg: (msg: string) => {
        const newMsg: ChatItem = {
          id: Date.now() + Math.random(),
          user_id: '1',
          sender_type: 'user',
          message: msg,
        }
        set({ list: [...get().list, newMsg] })
      },

      receive: (msg: ReceiveMsg) => {
        const list = [...get().list]
        if (msg.sender_type === 'user') {
          list.push(msg as ChatItem)
        } else if (msg.sender_type === 'chat') {
          const lastItem = list[list.length - 1]
          if (
            lastItem?.sender_type === 'chat' &&
            lastItem?.temporary === msg.temporary
          ) {
            lastItem.message += msg.data?.message ?? ''
          } else {
            list.push({
              ...msg.data,
              sender_type: 'chat',
              temporary: msg.temporary,
              id: Date.now() + Math.random(),
              user_id: '1',
              message: msg.data?.message ?? '',
            })
          }
        }
        set({ list })
      },

      clear: () => {
        set({ list: [] })
      },
    }),
    {
      name: 'chat-store',
      // partialize: (state) => ({ ...state, list: [] }), // 可自定义持久化字段
    }
  )
)

export default useStore
